import os
import hashlib
import shutil
from PIL import Image

def get_image_orientation(image_path):
    """判断图片的方向：横屏或竖屏"""
    with Image.open(image_path) as img:
        width, height = img.size
        return 'horizontal' if width > height else 'vertical'

def generate_md5(image_path):
    """生成图片的 MD5 值"""
    md5_hash = hashlib.md5()
    with open(image_path, 'rb') as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            md5_hash.update(byte_block)
    return md5_hash.hexdigest()

def classify_and_move_images(input_folder, horizontal_folder, vertical_folder):
    """根据图片方向分类并重命名图片，移动到不同文件夹"""
    if not os.path.exists(horizontal_folder):
        os.makedirs(horizontal_folder)
    
    if not os.path.exists(vertical_folder):
        os.makedirs(vertical_folder)

    # 获取输入文件夹下所有图片文件
    for filename in os.listdir(input_folder):
        image_path = os.path.join(input_folder, filename)
        
        if os.path.isfile(image_path):
            # 获取文件扩展名（后缀）
            file_extension = os.path.splitext(filename)[1]
            
            orientation = get_image_orientation(image_path)
            md5_value = generate_md5(image_path)
            
            # 根据方向决定新的文件名和保存路径
            if orientation == 'horizontal':
                new_filename = f'horizontal_{md5_value}{file_extension}'
                new_image_path = os.path.join(horizontal_folder, new_filename)
            else:
                new_filename = f'vertical_{md5_value}{file_extension}'
                new_image_path = os.path.join(vertical_folder, new_filename)
            
            try:
                # 打开图像并处理模式转换
                with Image.open(image_path) as img:
                    if img.mode == 'RGBA':
                        # 如果是 RGBA 图像，转换为 RGB 模式
                        img = img.convert('RGB')
                
                # 移动文件到对应的文件夹，并重命名
                shutil.move(image_path, new_image_path)
                print(f"图片 {filename} 已移动并重命名为 {new_filename} 到 {new_image_path}")
            
            except Exception as e:
                print(f"无法处理文件 {filename}: {e}")

# 示例：将 'input_images' 文件夹中的图片分类保存到 'horizontal_images' 和 'vertical_images' 文件夹
input_folder = './未作分类'
horizontal_folder = './horizontal_images'
vertical_folder = './vertical_images'

classify_and_move_images(input_folder, horizontal_folder, vertical_folder)
